package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/jmoiron/sqlx"
)

type QueryLog struct {
	Time         time.Time `json:"start_time"`
	User         string    `json:"user"`
	Host         string    `json:"host"`
	Ip           string    `json:"ip"`
	Pid          int       `json:"pid"`
	QueryTime    float64   `json:"query_time"`
	LockTime     float64   `json:"lock_time"`
	RowsSent     int       `json:"rows_sent"`
	RowsExamined int       `json:"rows_examined"`
	Sql          string    `json:"sql"`
}

func transfer(raw []byte) []QueryLog {
	reg := regexp.MustCompile("# +(.*)[\n\r]+# +(.*)[\n\r]+# +(.*)[\n\r]+")
	text := string(raw)
	sqls := reg.Split(text, -1)
	metas := reg.FindAllString(text, -1)
	timeReg := regexp.MustCompile("# +Time: +(.*)[\n\r]+")
	addrReg := regexp.MustCompile("# +User@Host: +(.*?)\\[(.*)\\] +@ (.*) +\\[(.*)\\] +Id: +(\\d+)[\n\r]+")
	perfReg := regexp.MustCompile("# +Query_time: +(.*) +Lock_time: +(.*) +Rows_sent: +(\\d+) +Rows_examined: +(\\d+)[\n\r]+")
	timestampReg := regexp.MustCompile("SET timestamp=\\d+;\n")
	logs := []QueryLog{}
	// fmt.Println(len(metas), len(sqls))

	for i := 0; i < len(metas); i++ {
		block := metas[i]
		// [, time]
		arr1 := timeReg.FindAllStringSubmatch(block, -1)
		// [,, user, host, ip, id]
		arr2 := addrReg.FindAllStringSubmatch(block, -1)
		// [, query_time, lock_time, rows_sent, rows_examined]
		arr3 := perfReg.FindAllStringSubmatch(block, -1)

		pid, _ := strconv.Atoi(arr2[0][5])
		queryTime, err := strconv.ParseFloat(strings.Trim(arr3[0][1], " "), 64)
		if err != nil {
			panic(err)
		}
		lockTime, _ := strconv.ParseFloat(arr3[0][2], 64)
		rowsSent, _ := strconv.Atoi(arr3[0][3])
		rowsExamined, _ := strconv.Atoi(arr3[0][3])

		timeStr := strings.Trim(arr1[0][1], "\r")
		startTime, _ := time.ParseInLocation("2006-01-02T15:04:05.000000Z", timeStr, time.UTC)
		sqlText := timestampReg.ReplaceAllString(sqls[i+1], "")

		log := QueryLog{
			Time:         startTime,
			User:         arr2[0][2],
			Host:         arr2[0][3],
			Ip:           arr2[0][4],
			Pid:          pid,
			QueryTime:    queryTime,
			LockTime:     lockTime,
			RowsSent:     rowsSent,
			RowsExamined: rowsExamined,
			Sql:          sqlText,
		}
		logs = append(logs, log)
	}

	return logs
}

var DB *sqlx.DB

func getDB(dsn string, table string) *sqlx.DB {
	if DB != nil {
		return DB
	}

	db, err := sqlx.Open("mysql", dsn)

	if err != nil {
		panic(err)
	}

	if _, err := db.Query(fmt.Sprintf("desc %s", table)); err != nil {
		_, err := db.Exec(fmt.Sprintf(`CREATE TABLE %s (
			id int(11) UNSIGNED NOT NULL AUTO_INCREMENT,
			ip varchar(50) NOT NULL DEFAULT '',
			host varchar(50) NOT NULL DEFAULT '',
			user varchar(255) NOT NULL DEFAULT '',
			pid int NOT NULL DEFAULT 0,
			start_time datetime NOT NULL,
			query_time decimal(11, 6) NOT NULL DEFAULT 0,
			lock_time decimal(11, 6) NOT NULL DEFAULT 0,
			rows_sent int(11) NOT NULL DEFAULT 0,
			rows_examined int(11) NOT NULL DEFAULT 0,
			sql_text longtext NOT NULL,
			createAt datetime NOT NULL DEFAULT CURRENT_TIMESTAMP(),
			updateAt datetime NOT NULL DEFAULT CURRENT_TIMESTAMP() ON UPDATE CURRENT_TIMESTAMP(),
			PRIMARY KEY (id))`, table))
		if err != nil {
			panic(err)
		}
	}

	DB = db
	return db
}

func insert(table string, logs [][]interface{}) int {
	if len(logs) == 0 {
		return 0
	}

	sql := fmt.Sprintf("insert into %s(ip, host, user, pid, start_time, query_time, lock_time, rows_sent, rows_examined, sql_text) values", table)
	values := []string{}
	args := []interface{}{}
	for i := 0; i < len(logs); i++ {
		log := logs[i]
		values = append(values, "(?,?,?,?,?,?,?,?,?,?)")
		args = append(args, log[3], log[2], log[1], log[4], log[0], log[5], log[6], log[7], log[8], log[9])
	}

	sql += strings.Join(values, ",")

	stmt, err := DB.Prepare(sql)

	if err != nil {
		panic(err)
	}

	res, err := stmt.Exec(args...)

	if err != nil {
		panic(err)
	}

	c, _ := res.RowsAffected()
	return int(c)
}

func saveToJson(logs [][]interface{}, file string) {
	json, _ := json.Marshal(logs)
	ioutil.WriteFile(file, json, 0777)
}

func saveToDb(logs [][]interface{}, table string) {
	var lastTime time.Time

	err := DB.Get(&lastTime, fmt.Sprintf("select start_time as lastTime from %s order by id desc limit 1", table))

	if err != nil && err != sql.ErrNoRows {
		panic(err)
	}

	i := 1
	count := len(logs)
	chunkSize := 50
	isNext := false

	for i < count {
		log := logs[i]

		if isNext == false {
			var startTime time.Time
			if v, ok := log[0].(time.Time); ok {
				startTime = v
			}

			if startTime.After(lastTime) {
				fmt.Println("Start time:", startTime)
				isNext = true
			} else {
				i++
				continue
			}
		}

		offset := math.Min(float64(i+chunkSize), float64(count))
		chunks := logs[i:int(offset)]

		rows := insert(table, chunks)
		fmt.Println(i, rows)

		i += chunkSize
	}
}

func catch() {
	if err := recover(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

var jsonFile = flag.String("out", "", "Output to json file")
var slowTableName = flag.String("table", "slow_logs", "Table name")
var dsn = flag.String("dsn", "", "MySQL DSN <user>:<password>@tcp(<host>:<prot>)/<database>")

func main() {
	// defer catch()
	flag.Usage = func() {
		fmt.Fprintf(flag.CommandLine.Output(), "myslow-cli\nUsage: myslow [options] <Query-slow file>\nOptions:\n")
		flag.PrintDefaults()
	}

	flag.Parse()
	args := flag.Args()

	if len(args) == 0 || args[0] == "" {
		flag.Usage()
		os.Exit(0)
	}

	if *dsn != "" {
		getDB(*dsn, *slowTableName)
	}

	raw, err := ioutil.ReadFile(args[0])

	if err != nil {
		panic(err)
	}

	data := transfer(raw)

	logs := [][]interface{}{}
	logs = append(logs, []interface{}{
		"Time",
		"User",
		"Host",
		"Ip",
		"Pid",
		"QueryTime",
		"LockTime",
		"RowsSent",
		"RowsExamined",
		"Sql",
	})
	for i := 0; i < len(data); i++ {
		logs = append(logs, []interface{}{
			data[i].Time,
			data[i].User,
			data[i].Host,
			data[i].Ip,
			data[i].Pid,
			data[i].QueryTime,
			data[i].LockTime,
			data[i].RowsSent,
			data[i].RowsExamined,
			data[i].Sql,
		})
	}

	if *jsonFile != "" {
		saveToJson(logs, *jsonFile)
	}

	if DB != nil {
		saveToDb(logs, *slowTableName)
	}

	fmt.Println("Output Done !\nTotal:", len(logs))
}
